#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>

using namespace std;
using namespace cv;

// 颜色识别函数
string colorRecognition(Mat& hsv, RotatedRect& rotatedRect)
{
    // 获取矩形的中心坐标
    Point2f center = rotatedRect.center;
    // 在中心位置提取颜色
    Vec3b pixel = hsv.at<Vec3b>(center);
    // 判断颜色
    if (pixel[0] > 100 && pixel[0] < 120) // 判断为红色
    {
        return "Red";
    }
    else if (pixel[0] > 10 && pixel[0] < 250) // 判断为蓝色
    {
        return "Blue";
    }
    return "";
}

// 计算角度函数
float calculateAngle(Point2f& prevCenter, Point2f& currentCenter)
{
    float dx = currentCenter.x - prevCenter.x;
    float dy = currentCenter.y - prevCenter.y;
    return atan2(dy, dx) * 180 / CV_PI;
}

int main()
{
    // 打开视频文件
    VideoCapture capture("/home/lwj/opencv.21/Buff09(1).MP4");
    if (!capture.isOpened())
    {
        cout << "Failed to open file." << endl;
        return -1;
    }
    
    Mat frame, hsv, binaryFrame, binary;
    Mat drawFrame;
    Mat kernel = getStructuringElement(MORPH_RECT, Size(3, 3));
    
    namedWindow("Draw Frame", WINDOW_NORMAL);
    
    Point2f prevCenter;
    auto prevTime = chrono::high_resolution_clock::now();
    
    while (true)
    {
        if (!capture.read(frame))
            break;

        // 对原图像进行腐蚀操作
        erode(frame, frame, kernel, Point(-1, -1), 3);

        cvtColor(frame, hsv, COLOR_BGR2HSV);

        // 设置颜色阈范围，提取二值图像
        inRange(hsv, Scalar(0, 73, 223), Scalar(179, 255, 255), binary);

        // 对二图像进行膨胀操作
        dilate(binary, binaryFrame, kernel, Point(-1, -1), 5);

        // 寻找廓
        vector<vector<Point>> contours;
        findContours(binaryFrame, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

        drawFrame = frame.clone();
        for (size_t i = 0; i < contours.size(); ++i)
        {
            RotatedRect rotatedRect = minAreaRect(contours[i]);
            Point2f vertices[4];
            rotatedRect.points(vertices);
            double area = contourArea(contours[i]);
            if(area<350)
                continue;
            for (int j = 0; j < 4; ++j)
            {
                line(drawFrame, vertices[j], vertices[(j + 1) % 4], Scalar(0, 255, 0), 2);
            }

            // 调用颜色识别函数，并根据面积大小进行标记
            string color = colorRecognition(hsv, rotatedRect);
            if (area < 2000)
            {
                putText(drawFrame, "R", rotatedRect.center, FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 0, 255), 2);
               
            }
            else
            {
                putText(drawFrame, color, rotatedRect.center, FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 0, 255), 2);
            }
        }

        imshow("Draw Frame", drawFrame);

        if (contours.size() > 0)
        {
            // 获取当前矩形的中心点坐标
            RotatedRect rotatedRect = minAreaRect(contours[0]);
            Point2f currentCenter = rotatedRect.center;

            // 计算角度变化速度
            auto currentTime = chrono::high_resolution_clock::now();
            chrono::duration<double> elapsedTime = currentTime - prevTime;
            double deltaTime = elapsedTime.count();

            // 调用计算角度函数，并计算角速度
            float angle = calculateAngle(prevCenter, currentCenter);
            float angularVelocity = angle / deltaTime;
            cout << "Angular Velocity: " << angularVelocity << " degrees/s" << endl;

            prevCenter = currentCenter;
            prevTime = currentTime;
        }

        if (waitKey(30) == 27)
            break;
    }

    capture.release();
    destroyAllWindows();

    return 0;
}
